#!/bin/bash

# DDoS攻击检查脚本
# 设置日志文件路径
LOG_FILE="/var/log/ddos_check.log"

# 默认阈值设置（用户未输入时使用）
DEFAULT_IP_LIMIT=50        # 单个IP的最大连接数默认值
DEFAULT_SYN_LIMIT=50       # SYN_RECV状态连接阈值默认值
DEFAULT_TOTAL_LIMIT=50     # 总连接数阈值默认值

# 实际使用的阈值（将由用户输入或默认值填充）
IP_CONNECTIONS_LIMIT=0
SYN_RECV_LIMIT=0
TOTAL_CONNECTION_LIMIT=0

# 获取当前时间戳
timestamp() {
    date +"%Y-%m-%d %T"
}

# 显示阈值设置帮助
show_threshold_help() {
    echo "阈值说明:"
    echo "1. 单个IP最大连接数: 超过此数值的IP将被标记为疑似异常"
    echo "2. SYN_RECV状态连接阈值: 超过此数值可能存在SYN洪水攻击"
    echo "3. 总连接数阈值: 服务器总连接数超过此数值将触发警告"
    echo "提示: 直接回车将使用默认值50"
    echo
}

# 配置阈值（支持用户输入或使用默认值）
configure_thresholds() {
    echo "==== 配置检查阈值 ===="
    show_threshold_help

    # 读取单个IP连接数阈值
    read -p "请输入单个IP最大连接数阈值 (默认: $DEFAULT_IP_LIMIT): " ip_limit
    if [[ -z "$ip_limit" || ! "$ip_limit" =~ ^[0-9]+$ ]]; then
        IP_CONNECTIONS_LIMIT=$DEFAULT_IP_LIMIT
        echo "使用默认值: $DEFAULT_IP_LIMIT"
    else
        IP_CONNECTIONS_LIMIT=$ip_limit
        echo "已设置: $ip_limit"
    fi

    # 读取SYN_RECV阈值
    read -p "请输入SYN_RECV状态连接阈值 (默认: $DEFAULT_SYN_LIMIT): " syn_limit
    if [[ -z "$syn_limit" || ! "$syn_limit" =~ ^[0-9]+$ ]]; then
        SYN_RECV_LIMIT=$DEFAULT_SYN_LIMIT
        echo "使用默认值: $DEFAULT_SYN_LIMIT"
    else
        SYN_RECV_LIMIT=$syn_limit
        echo "已设置: $syn_limit"
    fi

    # 读取总连接数阈值
    read -p "请输入总连接数阈值 (默认: $DEFAULT_TOTAL_LIMIT): " total_limit
    if [[ -z "$total_limit" || ! "$total_limit" =~ ^[0-9]+$ ]]; then
        TOTAL_CONNECTION_LIMIT=$DEFAULT_TOTAL_LIMIT
        echo "使用默认值: $DEFAULT_TOTAL_LIMIT"
    else
        TOTAL_CONNECTION_LIMIT=$total_limit
        echo "已设置: $total_limit"
    fi

    echo
    echo "==== 最终阈值配置 ===="
    echo "单个IP最大连接数: $IP_CONNECTIONS_LIMIT"
    echo "SYN_RECV状态连接阈值: $SYN_RECV_LIMIT"
    echo "总连接数阈值: $TOTAL_CONNECTION_LIMIT"
    echo "======================"
    echo
}

# 检查异常IP连接
check_ip_connections() {
    echo "[$(timestamp)] 正在检查高并发连接IP..."
    
    # 使用ss或netstat获取连接信息
    if command -v ss &> /dev/null; then
        cmd="ss -ntu | awk '{print \$6}' | cut -d: -f1 | sort | uniq -c | sort -nr"
    else
        cmd="netstat -ntu | awk '{print \$5}' | cut -d: -f1 | sort | uniq -c | sort -nr"
    fi

    eval $cmd | while read count ip; do
        if [[ "$count" -gt $IP_CONNECTIONS_LIMIT && "$ip" != "Address" ]]; then
            echo "[$(timestamp)] 警告! 疑似异常IP: $ip 当前连接数: $count (阈值: $IP_CONNECTIONS_LIMIT)" | tee -a $LOG_FILE
        fi
    done
}

# 检查SYN_RECV状态连接（SYN洪水攻击特征）
check_syn_flood() {
    echo "[$(timestamp)] 正在检查SYN_RECV状态连接..."
    
    if command -v ss &> /dev/null; then
        syn_count=$(ss -n state syn-recv | wc -l)
    else
        syn_count=$(netstat -n -t | grep SYN_RECV | wc -l)
    fi

    if [[ $syn_count -gt $SYN_RECV_LIMIT ]]; then
        echo "[$(timestamp)] 警告! SYN_RECV连接数异常: $syn_count (阈值: $SYN_RECV_LIMIT)" | tee -a $LOG_FILE
    fi
}

# 检查总连接数
check_total_connections() {
    echo "[$(timestamp)] 正在检查总连接数..."
    
    if command -v ss &> /dev/null; then
        total=$(ss -s | grep "TCP:" | awk '{print $2}')
    else
        total=$(netstat -ant | wc -l)
    fi

    if [[ $total -gt $TOTAL_CONNECTION_LIMIT ]]; then
        echo "[$(timestamp)] 警告! 总连接数异常: $total (阈值: $TOTAL_CONNECTION_LIMIT)" | tee -a $LOG_FILE
    fi
}

# 执行所有检查
main() {
    # 配置阈值
    configure_thresholds
    
    echo "==== 开始DDoS检查 ====" | tee -a $LOG_FILE
    check_ip_connections
    check_syn_flood
    check_total_connections
    echo "==== 检查完成 ====" | tee -a $LOG_FILE
}

# 以root权限运行
if [[ $EUID -ne 0 ]]; then
    echo "错误: 需要root权限运行此脚本" >&2
    exit 1
fi

main